package org.dreamwork.jasmine2.restful;

import org.dreamwork.misc.Base64;
import org.dreamwork.secure.IKeyFetcher;
import org.dreamwork.util.StringUtil;

import javax.servlet.ServletContext;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.URL;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.WeakHashMap;

/**
 * Created by IntelliJ IDEA.
 * User: seth.yang
 * Date: 2014/12/14
 * Time: 1:50
 */
public class RestfulAPIResourceKeyFetcher implements IKeyFetcher {
    public static final String KEY_RESTFUL_KEY_FETCHER = "restful-key-location";
    private static final WeakHashMap<String, PrivateKey> privateKeyCaches = new WeakHashMap<String, PrivateKey> ();
    private static final WeakHashMap<String, PublicKey> publicKeyCaches = new WeakHashMap<String, PublicKey> ();
    private static final WeakHashMap<String, X509Certificate> certificateCaches = new WeakHashMap<String, X509Certificate> ();

    private ServletContext application;

    public RestfulAPIResourceKeyFetcher (ServletContext application) {
        this.application = application;
    }

    @Override
    public PrivateKey getPrivateKey (String issuer) throws NoSuchAlgorithmException, InvalidKeySpecException {
        if (privateKeyCaches.containsKey (issuer))
            return privateKeyCaches.get (issuer);

        String resourceRoot = application.getInitParameter (KEY_RESTFUL_KEY_FETCHER);
        if (StringUtil.isEmpty (resourceRoot))
            throw new RuntimeException ("Missing restful-key-location config.");

        String location = resourceRoot + '/' + issuer + "/private-key.pem";
        InputStream in;
        try {
            URL url = application.getResource (location);
            if (url != null) {
                in = url.openStream ();
            } else {
                url = new URL (location);
                in = url.openStream ();
            }
            if (in != null) {
                PrivateKey key = readPrivateKey (in);
                privateKeyCaches.put (issuer, key);
                return key;
            }

            throw new InvalidKeySpecException ();
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }
    }

    @Override
    public PublicKey getPublicKey (String issuer) throws NoSuchAlgorithmException, InvalidKeySpecException {
        if (publicKeyCaches.containsKey (issuer))
            return publicKeyCaches.get (issuer);

        String resourceRoot = application.getInitParameter (KEY_RESTFUL_KEY_FETCHER);
        if (StringUtil.isEmpty (resourceRoot))
            throw new RuntimeException ("Missing restful-key-location config.");

        String location = resourceRoot + '/' + issuer + "/public-key.pem";
        InputStream in ;
        try {
            URL url = application.getResource (location);
            if (url != null) {
                in = url.openStream ();
            } else {
                url = new URL (location);
                in = url.openStream ();
            }
            if (in != null) {
                PublicKey kek = readPublicKey (in);
                publicKeyCaches.put (issuer, kek);
                return kek;
            }

            throw new InvalidKeySpecException ();
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }
    }

    @Override
    public X509Certificate getCertificate (String issuer) throws CertificateException {
        if (certificateCaches.containsKey (issuer))
            return certificateCaches.get (issuer);

        String resourceRoot = application.getInitParameter (KEY_RESTFUL_KEY_FETCHER);
        if (StringUtil.isEmpty (resourceRoot))
            throw new RuntimeException ("Missing restful-key-location config.");

        String location = resourceRoot + '/' + issuer + "/certificate.der";
        InputStream in;
        try {
            URL url = application.getResource (location);
            if (url != null) {
                in = url.openStream ();
            } else {
                url = new URL (location);
                in = url.openStream ();
            }
            if (in != null) {
                X509Certificate cert = readCertificate (in);
                certificateCaches.put (issuer, cert);
                return cert;
            }

            throw new InvalidKeySpecException ();
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }
    }

    private PrivateKey readPrivateKey (InputStream in) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
        StringBuilder builder = new StringBuilder ();
        BufferedReader reader = new BufferedReader (new InputStreamReader (in));
        String line;
        try {
            while ((line = reader.readLine ()) != null) {
                if (!line.startsWith ("--") && !StringUtil.isEmpty (line))
                    builder.append (line);
            }
        } finally {
            in.close ();
        }
        byte[] buff = Base64.decode (builder.toString ().getBytes ("utf-8"));
        if (buff != null) {
            KeyFactory factory = KeyFactory.getInstance ("RSA");
            KeySpec spec = new PKCS8EncodedKeySpec (buff);
            return factory.generatePrivate (spec);
        } else {
            throw new InvalidKeySpecException ();
        }
    }

    private PublicKey readPublicKey (InputStream in) throws IOException, InvalidKeySpecException, NoSuchAlgorithmException {
        try {
            BufferedReader reader = new BufferedReader (new InputStreamReader (in));
            String line;
            StringBuilder builder = new StringBuilder ();
            while ((line = reader.readLine ()) != null) {
                if (!line.startsWith ("--"))
                    builder.append (line);
            }

            byte[] data = Base64.decode (builder.toString ().trim ());
            X509EncodedKeySpec spec = new X509EncodedKeySpec (data);
            KeyFactory factory = KeyFactory.getInstance ("rsa");
            return factory.generatePublic (spec);
        } finally {
            in.close ();
        }
    }

    private X509Certificate readCertificate (InputStream in) throws CertificateException, IOException {
        try {
            CertificateFactory factory = CertificateFactory.getInstance ("X.509");
            return (X509Certificate) factory.generateCertificate (in);
        } finally {
            in.close ();
        }

    }
}